import { defineWrappedResponseHandler } from "#server/utils/handler";
import { R } from "#server/utils/response";
import { getAgentBySlug } from "#server/service/agent/agent";
import { resolveAgentIdentity } from "#server/service/agent/identity";
import { createTask, sendTask, getTask, cancelTask } from "#server/service/a2a/service";
import { DEFAULT_A2A_CONFIG, JSONRPC_ERROR_CODES, A2A_METHODS } from "#server/service/a2a/types";
import type {
  JSONRPCRequest,
  JSONRPCResponse,
  JSONRPCError,
  A2ATaskSendParams,
  A2ATaskGetParams,
  A2ATaskCancelParams,
  A2AInvocationContext,
  A2AMessage,
} from "#server/service/a2a/types";

function rpcError(id: string | number | null, code: number, message: string, data?: unknown): JSONRPCResponse {
  const error: JSONRPCError = { code, message };
  if (data !== undefined) error.data = data;
  return { jsonrpc: "2.0", id, error };
}

function rpcResult<T>(id: string | number | null, result: T): JSONRPCResponse {
  return { jsonrpc: "2.0", id, result };
}

export default defineWrappedResponseHandler(async (event) => {
  const agentSlug = getRouterParam(event, "agentSlug");
  if (!agentSlug) {
    return rpcError(null, JSONRPC_ERROR_CODES.INVALID_PARAMS, "missing agentSlug");
  }

  const agent = await getAgentBySlug(agentSlug);
  if (!agent) {
    return rpcError(null, JSONRPC_ERROR_CODES.AGENT_NOT_CALLABLE, `agent not found: ${agentSlug}`);
  }

  const body = await readBody<JSONRPCRequest>(event);
  if (!body || body.jsonrpc !== "2.0" || !body.method) {
    return rpcError(body?.id ?? null, JSONRPC_ERROR_CODES.INVALID_REQUEST, "invalid JSON-RPC 2.0 request");
  }

  const isInternal = event.node.req.headers["x-a2a-internal"] === "true";
  const internalCallerAgentId = event.node.req.headers["x-a2a-caller-agent-id"];
  const internalRecursionDepth = event.node.req.headers["x-a2a-recursion-depth"];

  let userId: number | null;
  let tempToken: string | null;
  if (isInternal) {
    userId = null;
    tempToken = null;
  } else {
    const identity = await resolveAgentIdentity(event);
    userId = identity.userId;
    tempToken = identity.tempToken;
  }

  const callerAgentId = isInternal
    ? (internalCallerAgentId ? Number(internalCallerAgentId) : null)
    : null;
  const recursionDepth = isInternal
    ? (internalRecursionDepth ? Number(internalRecursionDepth) : 0)
    : 0;

  const { id } = body;

  switch (body.method) {
    case A2A_METHODS.TASKS_SEND: {
      const params = body.params as A2ATaskSendParams;
      if (!params?.message || !params?.id) {
        return rpcError(id, JSONRPC_ERROR_CODES.INVALID_PARAMS, "missing required params: id, message");
      }

      const { task } = await createTask({
        calleeAgentSlug: agentSlug,
        message: params.message as A2AMessage,
        callerAgentId,
        userId,
        tempToken,
        existingSessionId: params.sessionId,
      });

      const ctx: A2AInvocationContext = {
        callerAgentId,
        calleeAgentId: agent.id,
        calleeAgentSlug: agentSlug,
        userId,
        sessionId: task.sessionId,
        recursionDepth,
        maxRecursionDepth: DEFAULT_A2A_CONFIG.maxRecursionDepth,
        timeoutMs: DEFAULT_A2A_CONFIG.timeoutMs,
        maxOutputTokens: DEFAULT_A2A_CONFIG.maxOutputTokens,
      };

      const completedTask = await sendTask(task.id, ctx);
      return rpcResult(id, completedTask);
    }

    case A2A_METHODS.TASKS_GET: {
      const params = body.params as A2ATaskGetParams;
      if (!params?.id) {
        return rpcError(id, JSONRPC_ERROR_CODES.INVALID_PARAMS, "missing required param: id");
      }
      const task = await getTask(params.id);
      if (!task) {
        return rpcError(id, JSONRPC_ERROR_CODES.TASK_NOT_FOUND, `task not found: ${params.id}`);
      }
      return rpcResult(id, task);
    }

    case A2A_METHODS.TASKS_CANCEL: {
      const params = body.params as A2ATaskCancelParams;
      if (!params?.id) {
        return rpcError(id, JSONRPC_ERROR_CODES.INVALID_PARAMS, "missing required param: id");
      }
      const canceled = await cancelTask(params.id);
      if (!canceled) {
        return rpcError(id, JSONRPC_ERROR_CODES.TASK_NOT_CANCELABLE, `task not cancelable: ${params.id}`);
      }
      return rpcResult(id, { id: params.id, canceled: true });
    }

    case A2A_METHODS.TASKS_SEND_SUBSCRIBE:
    case A2A_METHODS.TASKS_SUBSCRIBE:
      return rpcError(id, JSONRPC_ERROR_CODES.NOT_IMPLEMENTED, "streaming not implemented yet");

    default:
      return rpcError(id, JSONRPC_ERROR_CODES.METHOD_NOT_FOUND, `method not found: ${body.method}`);
  }
});